{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Run upon export from spreadsheet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total entries: 590\n",
      "Total labeled entries: 577\n",
      "Split sizes. Train: 0; Valid: 0; Test: 0\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "\n",
    "from astroquery.mast import Catalogs\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "\n",
    "tces_file = '/mnt/tess/labels/s33_cam2ccd14_sample.csv'\n",
    "ext_data_file = '/mnt/tess/labels/ext_mast_data.csv'\n",
    "labels_file = '/mnt/tess/labels/labels_ext_mission_test.csv'\n",
    "splits_file = '/mnt/tess/labels/splits_v3.csv'\n",
    "\n",
    "\n",
    "tce_table = pd.read_csv(tces_file, header=0, low_memory=False)\n",
    "tce_table['tic_id'] = tce_table['star_tic']\n",
    "tce_table['Duration'] = tce_table['planet_tdur']\n",
    "tce_table['Period'] = tce_table['planet_period']\n",
    "tce_table['RA'] = tce_table['star_ra']\n",
    "tce_table['Sectors'] = tce_table['sector_id'].apply(lambda v: len(v.split(' ')))\n",
    "tce_table['Transit_Depth'] = tce_table['planet_depth']\n",
    "tce_table['Dec'] = tce_table['star_dec']\n",
    "tce_table['teff'] = tce_table['star_teff']\n",
    "tce_table['SN'] = tce_table['snr']\n",
    "tce_table['Qingress'] = 0.0\n",
    "tce_table['Tmag'] = tce_table['star_tmag']\n",
    "tce_table['logg'] = tce_table['star_logg']\n",
    "tce_table['Epoc'] = tce_table['planet_epoch']\n",
    "tce_table = tce_table.set_index('tic_id')\n",
    "tce_table = tce_table.drop(columns=['Unnamed: 0'])\n",
    "tce_table['Duration'] /= 24.0\n",
    "\n",
    "# Drop some common invalid examples.\n",
    "# Orbits falling inside the star\n",
    "tce_table = tce_table[~tce_table.Ilabel]\n",
    "# Excessively large durations\n",
    "tce_table = tce_table[tce_table.Duration < 0.9 * tce_table.Period]\n",
    "\n",
    "joined_table = tce_table\n",
    "\n",
    "ext_table = pd.read_csv(ext_data_file, header=0, low_memory=False).set_index('tic_id')\n",
    "joined_table = joined_table.join(ext_table, on='tic_id', how='left')\n",
    "\n",
    "joined_table = joined_table[\n",
    "    joined_table['objType'].isnull()\n",
    "    | (joined_table['objType'] == 'STAR')\n",
    "]\n",
    "\n",
    "joined_table = joined_table.reset_index()[[\n",
    "    'tic_id', 'RA', 'Dec', 'Tmag', 'Epoc', 'Period', 'Duration',\n",
    "    'Transit_Depth', 'Sectors', 'star_rad', 'star_mass', 'teff',\n",
    "    'logg', 'SN', 'Qingress'\n",
    "]]\n",
    "\n",
    "\n",
    "labels_table = pd.read_csv(labels_file, header=0, low_memory=False)\n",
    "disps = ['E', 'J', 'N', 'S', 'B']\n",
    "users = ['av', 'md', 'ch', 'as', 'mk', 'et']\n",
    "\n",
    "for d in disps:\n",
    "    labels_table[f'disp_{d}'] = 0\n",
    "\n",
    "def set_labels(row):\n",
    "    a = ~row.isna()\n",
    "    if a['Final']:\n",
    "        row[f'disp_{row[\"Final\"]}'] = 1\n",
    "    else:\n",
    "        for user in users:\n",
    "#             # Override md's votes when they're 1-to-all against J\n",
    "#             if user == 'md' and row[user] in ('B', 'N'):\n",
    "#                 others = [\n",
    "#                     row[u] for u in users\n",
    "#                     if row[u] and u != 'md' and not(isinstance(row[u], float) and np.isnan(row[u]))]\n",
    "#                 if all(o == 'J' for o in others):\n",
    "#                     row[f'disp_J'] += 1\n",
    "#                     continue\n",
    "            if a[user] and row[user] and row[user] != 'U':\n",
    "                row[f'disp_{row[user]}'] += 1\n",
    "                        \n",
    "    return row\n",
    "\n",
    "labels_table['tic_id'] = labels_table['TIC ID']\n",
    "labels_table = labels_table.apply(set_labels, axis=1)\n",
    "\n",
    "labels_table = labels_table[['tic_id'] + [f'disp_{d}' for d in disps]]\n",
    "\n",
    "\n",
    "joined_table = joined_table.set_index('tic_id')\n",
    "labels_table = labels_table.set_index('tic_id')\n",
    "joined_table = joined_table.join(labels_table, on='tic_id', how='inner')\n",
    "print(f'Total entries: {len(joined_table)}')\n",
    "joined_table = joined_table[\n",
    "    sum(joined_table[f'disp_{d}'] for d in disps) > 0\n",
    "]\n",
    "print(f'Total labeled entries: {len(joined_table)}')\n",
    "\n",
    "\n",
    "all_table = joined_table\n",
    "splits_table = pd.read_csv(splits_file, header=0, low_memory=False)\n",
    "splits_table['tic_id'] = splits_table['TIC ID']\n",
    "splits_table = splits_table.set_index('tic_id')\n",
    "joined_table = joined_table.join(splits_table, on='tic_id', how='inner')\n",
    "\n",
    "t_train = joined_table[joined_table['Split'] == 'train']\n",
    "t_val = joined_table[joined_table['Split'] == 'val']\n",
    "t_test = joined_table[joined_table['Split'] == 'test']\n",
    "t_train = t_train.drop(columns=['Hemisphere', 'Seed randbetween(1, 100)', 'Split'])\n",
    "t_val = t_val.drop(columns=['Hemisphere', 'Seed randbetween(1, 100)', 'Split'])\n",
    "t_test = t_test.drop(columns=['Hemisphere', 'Seed randbetween(1, 100)', 'Split'])\n",
    "print(f'Split sizes. Train: {len(t_train)}; Valid: {len(t_val)}; Test: {len(t_test)}')\n",
    "\n",
    "\n",
    "# t_train.to_csv('/mnt/tess/astronet/tces-v6-train.csv')\n",
    "# t_val.to_csv('/mnt/tess/astronet/tces-v6-val.csv')\n",
    "# t_test.to_csv('/mnt/tess/astronet/tces-v6-test.csv')\n",
    "all_table.to_csv('/mnt/tess/astronet/tces-tmp-ext-all.csv')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pd.set_option('display.max_columns', None)\n",
    "t_train.sample(5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t_val.sample(5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "t_test.sample(5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Run once"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_tces_old():\n",
    "    tceold = pd.read_csv('/mnt/tess/astronet/tces.csv', header=0).set_index('tic_id')\n",
    "\n",
    "    # Only keep the max sectors read.\n",
    "    maxsect = tceold.groupby('tic_id')['Sectors'].max()\n",
    "    tceold = tceold.join(maxsect, on='tic_id', how='right', rsuffix='_max')\n",
    "    tceold = tceold[tceold.Sectors == tceold.Sectors_max]\n",
    "\n",
    "    # Then keep the max row ID.\n",
    "    maxrowid = tceold.groupby('tic_id')['row_id'].max()\n",
    "    tceold = tceold.join(maxrowid, on='tic_id', how='right', rsuffix='_max')\n",
    "    tceold = tceold[tceold.row_id == tceold.row_id_max]\n",
    "\n",
    "    return tceold\n",
    "\n",
    "def generate_tce_bls_instar():\n",
    "    tcenew = pd.read_csv('/mnt/tess/labels/tce_bls_instar.csv', header=0).set_index('tic_id')\n",
    "    tceold = load_tces_old()\n",
    "    tcenorth = pd.read_csv('/mnt/tess/labels/tce_north_instar.csv', header=0).set_index('tic_id')\n",
    "\n",
    "    # Copy from old data where it's missing from the new.\n",
    "    alltce = tcenew.join(tceold, how='outer', on='tic_id', rsuffix='_old')\n",
    "    alltce = alltce.set_index('tic_id')\n",
    "\n",
    "    alltce = alltce.drop(columns=['row_id'])\n",
    "\n",
    "    def fillna(df, col_name):\n",
    "        df.loc[df[col_name].isna(), col_name] = df.loc[df[col_name].isna(), col_name + '_old']\n",
    "\n",
    "    fillna(alltce, 'toi_id')\n",
    "    fillna(alltce, 'Disposition')\n",
    "    fillna(alltce, 'RA')\n",
    "    fillna(alltce, 'Dec')\n",
    "    fillna(alltce, 'Tmag')\n",
    "    fillna(alltce, 'Epoc')\n",
    "    fillna(alltce, 'Period')\n",
    "    fillna(alltce, 'Duration')\n",
    "    fillna(alltce, 'Transit_Depth')\n",
    "    fillna(alltce, 'Sectors')\n",
    "    fillna(alltce, 'camera')\n",
    "    fillna(alltce, 'ccd')\n",
    "    fillna(alltce, 'star_rad')\n",
    "    fillna(alltce, 'star_mass')\n",
    "    fillna(alltce, 'teff')\n",
    "    fillna(alltce, 'logg')\n",
    "    fillna(alltce, 'SN')\n",
    "    fillna(alltce, 'Qingress')\n",
    "\n",
    "    alltce = alltce.drop(columns=[c for c in alltce.columns if c.endswith('_old')])\n",
    "    \n",
    "    alltce = alltce.append(tcenorth)\n",
    "    \n",
    "    alltce['Ilabel'] = alltce['Ilabel'].fillna(False)\n",
    "\n",
    "    alltce.to_csv('/mnt/tess/labels/tce_bls_instar+old.csv')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
